"""
文档加载器
"""
import os
import importlib
from typing import List, Dict, Any, Optional
from pathlib import Path
from langchain.docstore.document import Document
from langchain_community.document_loaders import TextLoader

from ..config.constants import LOADER_DICT
from ..config.settings import settings


class DocumentLoader:
    """文档加载器"""
    
    def __init__(self):
        self.supported_extensions = self._get_supported_extensions()
    
    def _get_supported_extensions(self) -> List[str]:
        """获取支持的文件扩展名"""
        extensions = []
        for loader_name, exts in LOADER_DICT.items():
            extensions.extend(exts)
        return extensions
    
    def get_loader_class(self, file_extension: str) -> str:
        """根据文件扩展名获取加载器类名"""
        file_extension = file_extension.lower()
        for loader_name, extensions in LOADER_DICT.items():
            if file_extension in extensions:
                return loader_name
        raise ValueError(f"不支持的文件格式: {file_extension}")
    
    def get_loader(self, file_path: str, loader_kwargs: Dict[str, Any] = None) -> Any:
        """获取文档加载器实例"""
        loader_kwargs = loader_kwargs or {}
        file_extension = Path(file_path).suffix.lower()
        
        try:
            loader_name = self.get_loader_class(file_extension)
            
            # 特殊加载器处理
            if loader_name in [
                "RapidOCRPDFLoader",
                "RapidOCRLoader", 
                "RapidOCRDocLoader",
                "RapidOCRPPTLoader",
            ]:
                # 这里可以导入自定义的OCR加载器
                # 暂时使用UnstructuredFileLoader作为fallback
                loader_name = "UnstructuredFileLoader"
            
            # 导入加载器模块
            if loader_name == "UnstructuredFileLoader":
                module = importlib.import_module("langchain_community.document_loaders.unstructured")
            else:
                module = importlib.import_module("langchain_community.document_loaders")
            
            # 获取加载器类
            loader_class = getattr(module, loader_name)
            
            # 创建加载器实例
            if loader_name == "TextLoader":
                loader = loader_class(file_path, encoding="utf8", **loader_kwargs)
            else:
                loader = loader_class(file_path, **loader_kwargs)
            
            return loader
            
        except Exception as e:
            raise Exception(f"加载文档 {file_path} 时出错: {e}")
    
    def load_document(self, file_path: str, loader_kwargs: Dict[str, Any] = None) -> List[Document]:
        """加载单个文档"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        loader = self.get_loader(file_path, loader_kwargs)
        return loader.load()
    
    def load_documents(self, file_paths: List[str], loader_kwargs: Dict[str, Any] = None) -> List[Document]:
        """批量加载文档"""
        documents = []
        for file_path in file_paths:
            try:
                docs = self.load_document(file_path, loader_kwargs)
                documents.extend(docs)
            except Exception as e:
                print(f"加载文档 {file_path} 失败: {e}")
                continue
        return documents
    
    def is_supported(self, file_path: str) -> bool:
        """检查文件格式是否支持"""
        file_extension = Path(file_path).suffix.lower()
        return file_extension in self.supported_extensions
    
    def get_file_info(self, file_path: str) -> Dict[str, Any]:
        """获取文件信息"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        path = Path(file_path)
        return {
            "file_path": str(path),
            "file_name": path.name,
            "file_extension": path.suffix.lower(),
            "file_size": os.path.getsize(file_path),
            "is_supported": self.is_supported(file_path),
        }


class KnowledgeFile:
    """知识库文件类"""
    
    def __init__(self, file_path: str, knowledge_base_name: str = "default"):
        self.file_path = file_path
        self.kb_name = knowledge_base_name
        self.loader = DocumentLoader()
        self.docs = None
        self.splited_docs = None
        
        # 验证文件
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        if not self.loader.is_supported(file_path):
            raise ValueError(f"不支持的文件格式: {Path(file_path).suffix}")
    
    def file2docs(self, refresh: bool = False) -> List[Document]:
        """将文件转换为文档列表"""
        if self.docs is None or refresh:
            self.docs = self.loader.load_document(self.file_path)
        return self.docs
    
    def get_file_info(self) -> Dict[str, Any]:
        """获取文件信息"""
        return self.loader.get_file_info(self.file_path) 